from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve
from tqdm import tqdm

plt.rcParams['axes.unicode_minus'] = False
plt.rc('font', size=10)
plt.rc('font', family='SimHei')

v = 1
# b = 0.45 / (2 * np.pi)
theta_0 = 16 * 2 * np.pi
l_head = 3.41 - 0.275 * 2
l_body = 2.2 - 0.275 * 2
l_0 = 0.275
half_width = 0.15
label_list = [0, 1, 51, 101, 151, 201, 222]


# theta与t的关系
def pfun(y, t):
    global v, b
    return (-1) * v / (b * np.sqrt(1 + y ** 2))


def x(theta, r):
    return r * np.cos(theta)


def y(theta, r):
    return r * np.sin(theta)


# 计算速度向量的函数
def velocity_from_theta(theta, dtheta_dt, b):
    u_x = np.cos(theta) - theta * np.sin(theta)
    u_y = np.sin(theta) + theta * np.cos(theta)
    return b * dtheta_dt * np.array([u_x, u_y])


"""
求解极坐标系下圆和等距螺线的交点
"""
def find_intersections(r0, phi0, R, b, theta_guesses):
    def equations(p):
        theta = p[0]
        r = b * theta
        eq1 = r ** 2 + r0 ** 2 - 2 * r * r0 * np.cos(theta - phi0) - R ** 2
        return [eq1]

    solutions = []
    for guess in theta_guesses:
        theta_sol = fsolve(equations, guess)[0]
        r_sol = b * theta_sol
        # if theta_sol > 0 and abs(equations([theta_sol])[0]) < 1e-2:
        if theta_sol > 0:
            is_new = True
            for sol in solutions:
                if np.allclose([theta_sol], [sol[0]], atol=1e-2):
                    is_new = False
                    break
            if is_new:
                solutions.append((theta_sol, r_sol))
    solutions.sort()
    return solutions


# 修改后的find_next函数，现在返回位置和速度，并递归计算后续棍头
def find_next(prev_theta, prev_r, prev_pos_cart, prev_v, i=0):
    # 使用前一个棍头的极坐标作为圆心
    theta_guesses = [prev_theta + 0.001, prev_theta + 0.01, prev_theta + 0.1]  # 初始猜测角度
    intersections = []
    if i == 0:
        intersections = find_intersections(prev_r, prev_theta, l_head, b, theta_guesses)
    else:
        intersections = find_intersections(prev_r, prev_theta, l_body, b, theta_guesses)

    results = []
    if not intersections:
        return intersections
    theta_sol, r_sol = intersections[0]
    # 当前棍头的笛卡尔坐标
    pos_cart = np.array([x(theta_sol, r_sol), y(theta_sol, r_sol)])
    # 计算向量 delta_r = current - previous
    delta_r = pos_cart - prev_pos_cart

    # 计算 u向量
    u_vec = np.array([np.cos(theta_sol) - theta_sol * np.sin(theta_sol),
                      np.sin(theta_sol) + theta_sol * np.cos(theta_sol)])

    # 计算分母和分子
    denominator = b * np.dot(u_vec, delta_r)
    numerator = np.dot(prev_v, delta_r)

    if np.abs(denominator) < 1e-6:
        # 避免除零，跳过或处理
        dtheta_dt = 0
    else:
        dtheta_dt = numerator / denominator

    # 计算当前速度
    current_v = velocity_from_theta(theta_sol, dtheta_dt, b)

    # 存储结果：位置、速度、极角、极径
    results.append({
        'theta': theta_sol,
        'r': r_sol,
        'pos_cart': pos_cart,
        'velocity': current_v,
        'index': i
    })

    if i > 222 or theta_sol > soli[0]:
        return []
    else:
        # print(pos_cart)
        # if i in label_list:
            # print(f"i={i}时，位置：{x(prev_theta, prev_r), y(prev_theta, prev_r)}，速度：{(prev_v[0] ** 2 + prev_v[1] ** 2) ** 0.5}")
            # print(f"b * dtheta_dt:{b * dtheta_dt}")
        # 递归计算下一个棍头，传递当前棍头的参数
        next_results = find_next(theta_sol, r_sol, pos_cart, current_v, i + 1)
        results.extend(next_results)
    return results


# 主程序
t = np.arange(0, 420, 0.1)
# for b in (np.arange(0.4, 0.5, 0.01) / (2 * np.pi))[::-1]:
for b in tqdm((np.arange(0.46, 0.5, 0.01) / (2 * np.pi))):
    soli = odeint(pfun, theta_0, t)[:, 0]
    # 选择时间点，例如t_index=600（对应t=60s，因为步长0.1，索引600）
    for t_index in np.arange(0, 420, 0.5):
        # print(t_index)
        theta0 = soli[int(t_index * 10)]
        r0 = b * theta0
        dtheta0_dt = pfun(theta0, t[int(t_index * 10)])
        v0 = velocity_from_theta(theta0, dtheta0_dt, b)
        pos0_cart = np.array([x(theta0, r0), y(theta0, r0)])

        # 计算所有棍头的位置和速度
        all_heads = find_next(theta0, r0, pos0_cart, v0)
        # 绘制其他棍头
        collision_points = []  # 观察图像可得，如果要发生碰撞，一定会发生在龙头左上方的点处，即point_4，因此记录下它的坐标
        collision_side = []  # 而可能发生碰撞的边也会是在边框的左边
        for i in range(len(all_heads)):
            pos = all_heads[i]['pos_cart']
            if i > 0:
                prev_pos = all_heads[i - 1]['pos_cart']
            else:
                prev_pos = [pos0_cart[0], pos0_cart[1]]
            l = np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # 由b指向a的向量
            l_back = -1 * np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # l向量的相反向量
            l_len = np.linalg.norm(l)
            d = np.array([(pos[0] - prev_pos[0]) / l_len, (pos[1] - prev_pos[1]) / l_len])  # l向量平行的单位向量
            n_left = np.array([-d[1], d[0]])  # 线段方向逆时针转得到的法向量
            n_right = np.array([d[1], -d[0]])  # 线段方向顺时针转得到的法向量
            n_len = np.linalg.norm(n_left)
            point_1 = pos + l_0 * d + half_width * n_left
            point_2 = prev_pos - l_0 * d + half_width * n_left
            point_3 = pos + l_0 * d - half_width * n_left
            point_4 = prev_pos - l_0 * d - half_width * n_left
            if i == 0:
                collision_points = [point_4, point_3]
            else:
                collision_side.append([[point_1[0], point_1[1]], [point_2[0], point_2[1]]])


        def is_point_on_segment(point, seg_start, seg_end):
            # 解包坐标
            px, py = point
            x1, y1 = seg_start
            x2, y2 = seg_end

            # 计算向量叉积（判断点是否在直线上）
            cross_product = (px - x1) * (y2 - y1) - (py - y1) * (x2 - x1)
            # 若叉积非零（考虑浮点精度），则点不在直线上
            if abs(cross_product) > 0.01:
                return False

            # 检查点是否在线段的轴向包围盒内（包括端点）
            if min(x1, x2) <= px <= max(x1, x2) and min(y1, y2) <= py <= max(y1, y2):
                # print(f"时间{t_index}时最小叉积为{cross_product}")
                return True

            return False

        flag = False
        # 遍历 collision_side 检测碰撞
        for collision_point in collision_points:
            for i in range(len(collision_side)):
                seg_start = collision_side[i][0]
                seg_end = collision_side[i][1]

                if is_point_on_segment(collision_point, seg_start, seg_end):
                    print(f"点{collision_point}与线段{seg_start} - {seg_end}发生碰撞，此时时间为{t_index}，龙头的半径为{(pos0_cart[0] ** 2 + pos0_cart[1] ** 2) ** 0.5}，螺距为{b * 2 * np.pi}")
                    flag = True
                    break
            if flag == True:
                break
        if flag and (pos0_cart[0] ** 2 + pos0_cart[1] ** 2) ** 0.5 > 4.5:
            # 绘制图形
            plt.figure(figsize=(8, 8))
            # 绘制螺线
            theta_plot = np.linspace(0, soli[0], 1000)
            r_plot = b * theta_plot
            plt.plot(x(theta_plot, r_plot), y(theta_plot, r_plot), color='red', alpha=0.5)
            # 绘制第一个棍头
            plt.plot(pos0_cart[0], pos0_cart[1], 'ro', color='red')
            # 绘制其他棍头
            for i in range(len(all_heads)):
                pos = all_heads[i]['pos_cart']
                plt.plot(pos[0], pos[1], 'ro', color='red')
                if i > 0:
                    prev_pos = all_heads[i - 1]['pos_cart']
                else:
                    prev_pos = [pos0_cart[0], pos0_cart[1]]
                l = np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # 由b指向a的向量
                l_back = -1 * np.array([pos[0] - prev_pos[0], pos[1] - prev_pos[1]])  # l向量的相反向量
                l_len = np.linalg.norm(l)
                d = np.array([(pos[0] - prev_pos[0]) / l_len, (pos[1] - prev_pos[1]) / l_len])  # l向量平行的单位向量
                n_left = np.array([-d[1], d[0]])  # 线段方向逆时针转得到的法向量
                n_right = np.array([d[1], -d[0]])  # 线段方向顺时针转得到的法向量
                n_len = np.linalg.norm(n_left)
                point_1 = pos + l_0 * d + half_width * n_left
                point_2 = prev_pos - l_0 * d + half_width * n_left
                point_3 = pos + l_0 * d - half_width * n_left
                point_4 = prev_pos - l_0 * d - half_width * n_left

                # 绘制图形
                plt.plot([point_1[0], point_3[0]], [point_1[1], point_3[1]], color="blue", alpha=0.5)
                plt.plot([point_2[0], point_4[0]], [point_2[1], point_4[1]], color="blue", alpha=0.5)
                plt.plot([point_3[0], point_4[0]], [point_3[1], point_4[1]], color="blue", alpha=0.5)
                plt.plot([point_1[0], point_2[0]], [point_1[1], point_2[1]], color="black", alpha=0.5)
                if i == 0:
                    plt.plot(point_3[0], point_3[1], "^", color="black", markersize=4)
                    plt.plot(point_4[0], point_4[1], "^", color="black", markersize=4)
            plt.xlabel('X')
            plt.ylabel('Y')
            plt.title('把手位置')
            plt.axis('equal')
            plt.grid(True)
            plt.show()
            break
